import numpy as np
import acado
import time
import matplotlib.pyplot as plt
from numpy import sin, cos,exp

import time
import rospy

N=20 #timestep
NU=4
NOD=8
NX=16
NY=16
NYN=4

NW=16
NWN=12

x0=np.zeros(NX).reshape(1,NX)
od=np.zeros(NOD)
y=np.zeros(NY)
u=np.zeros(NU)
Q=np.eye(NW)


#straight line 
#Experimental group
Q=np.diag([0,0,0,1,1,1,1,100,10,10,100,10,1,1,1,5e5])
#Control group
#Q=np.diag([0,0,0,1,1,1,1,100,10,10,1,1000,1,1,1,5e5])

#circle
#Experiment group
#Q=np.diag([0,0,0,1,1,1,1,100,100,100,100,100000,1,1,1,5e5])
#Control group
#Q=np.diag([0,0,0,1,1,1,1,100,100,100,1,100000,1,1,1,5e5])



W=np.transpose(np.tile(Q,N))


WN=Q[:-4,:-4]

y=np.zeros([N,NY])


x0[0,0]=-1.5 #px
x0[0,1]=0.0#py
x0[0,2]=0.0 #pz
x0[0,3]=1.0#qw
x0[0,4]=0.0#qx
x0[0,5]=0.0#qy
x0[0,6]=0.0#qz
x0[0,7]=0.0#vx
x0[0,8]=0.0#vy
x0[0,9]=0.0#vz
x0[0,10]=1.0#pxx
x0[0,11]=0.0#pxy
x0[0,12]=0.0#pxz
x0[0,13]=1.0#pyy
x0[0,14]=0.0#pyz
x0[0,15]=1.0#pzz

X=np.tile(x0.reshape(NX,1),N+1).T
OD=np.ones([N+1,NOD])
u=np.zeros([N,NU])

xlog=[]
tracelog=[]

r=3.3
cy=-r
cx=0

for i in range(1800):
    
    #straight line
    
    
    OD[:,0]=i*0.05
    OD[:,1]=sin(i*0.01)*0.+0.001
    OD[:,2]=1.0
    OD[:,3]=100.
    OD[:,4]=1.
    OD[:,5:]=0.1
    
    v=(1+sin(i*0.01))*0.01
    OD[:,0]+=v
    
    #print("iter: "+str(i))
    
    #circle
    theta=i*0.005
    #OD[:,0]=cx+r*sin(theta)
    #OD[:,1]=cy+r*cos(theta)
    
    X,u=acado.mpc(0,1,x0,X,u,OD,y,W,WN,0)
    
    x0=X[1,:].reshape(1,NX)
    #print(x0[0,7:10])# velocity
    

    
    dist=np.sqrt(np.sum((x0[0,0:3]-OD[-1,0:3])**2))
    
    print("dist: "+str(dist))
    
    P=np.array([[x0[0,10],x0[0,11],x0[0,12]],
                [x0[0,11],x0[0,13],x0[0,14]],
                [x0[0,12],x0[0,14],x0[0,15]]])
    
    
    #tracelog.append(x0[0,10]+x0[0,13]+x0[0,15])
    #tracelog.append(np.linalg.det(P))
    tracelog.append(x0[0,10]*x0[0,13]-x0[0,11]*x0[0,11])
    print(tracelog[-1])
    
    xlog.append(x0)
    xlog_arr=np.vstack(xlog)
    

    
    plt.subplot(211)
    plt.plot(xlog_arr[:,0],xlog_arr[:,1],"-r")
    plt.gcf().canvas.mpl_connect('key_release_event',
                                             lambda event: [exit(
                                                 0) if event.key == 'escape' else None])
   
    
    plt.scatter(OD[0,0],OD[0,1])
    
    plt.subplot(212)
    plt.title('traceP')
    plt.plot(tracelog)
    #plt.ylim([-0.01,0.2])
    
    if i%10==0:
        

        plt.pause(0.0001)

t1=time.time()

print(t1-t0)


plt.subplot(221)
plt.plot(xlog_arr[:,0])
plt.subplot(222)
plt.plot(xlog_arr[:,1])
plt.subplot(223)
plt.plot(tracelog)

plt.title('tracking result')



    
 
    
plt.grid()
plt.subplot(224)
plt.plot(xlog_arr[:,0],xlog_arr[:,1])
plt.title('determinat of covariance')

plt.show()

